import os
import pandas as pd
import re
from collections import defaultdict

# 定义输入和输出目录
input_dir = r'F:\Personal\data'
output_dir = r'F:\Personal\data\merged'

# 确保输出目录存在
os.makedirs(output_dir, exist_ok=True)

# 正则表达式用于匹配{name}_{YYYYMMDD}.csv
pattern = re.compile(r'^(.*?)_\d{8}\.csv$')

# 字典用于存储不同name对应的DataFrame列表
dataframes = defaultdict(list)

# 遍历输入目录下的所有文件
for filename in os.listdir(input_dir):
    # 跳过以stk_factor_开头的文件
    # if filename.startswith('stk_factor_') and filename.endswith('.csv'):
    #     continue

    # 跳过以stk_factor_开头的文件
    if not (filename.startswith('stk_factor_') and filename.endswith('.csv')):
        continue

    # 检查文件名是否匹配模式
    match = pattern.match(filename)
    if match:
        # 提取name
        name = match.group(1)

        # 构建文件的完整路径
        file_path = os.path.join(input_dir, filename)

        # 读取CSV文件到DataFrame
        df = pd.read_csv(file_path)

        # 添加到对应name的DataFrame列表
        dataframes[name].append(df)

# 对每个name的DataFrame列表进行合并并保存
for name, dfs in dataframes.items():
    # 合并所有DataFrame
    merged_df = pd.concat(dfs, ignore_index=True)

    # 定义输出文件名
    output_filename = f'{name}.csv'
    output_filepath = os.path.join(output_dir, output_filename)

    # 将合并后的DataFrame保存到CSV文件
    merged_df.to_csv(output_filepath, index=False)

    print(f"已合并文件到 {output_filepath}")
